#!/usr/bin/python3
# -*- coding: utf-8 -*-

# *************************************************************************
#    > File Name: server.py
#    > Author: xlzh
#    > Mail: xiaolongzhang2015@163.com
#    > Created Time: 2025年01月14日 星期二 14时31分29秒
# *************************************************************************

from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from tortoise.contrib.fastapi import register_tortoise
from models import PublicModel
from settings import *
import os.path
from typing import Optional
import uvicorn
import gzip


class Read(object):
    def __init__(self, title: str, seq: str, comment: str, quality: str):
        self.title = title
        self.seq = seq
        self.comment = comment
        self.quality = quality

    def read_length(self):
        return len(self.title) + len(self.seq) + len(self.comment) + len(self.quality)

    def to_string(self):
        return ''.join([self.title, self.seq, self.comment, self.quality])


class FastqIterator:
    def __init__(self, file_path):
        self.file_hd = None
        self.__open_file__(file_path)

    def __open_file__(self, file_path):
        if not os.path.exists(file_path):
            raise FileExistsError(f"Error: file of {file_path} is not exist!")

        if file_path.endswith('.gz'):
            self.file_hd = gzip.open(file_path, 'rt')
        else:
            self.file_hd = open(file_path, 'r')

    def get_read(self) -> Optional[Read]:
        group = []

        for idx in range(4):
            read_line = self.file_hd.readline()
            if not read_line:
                return None
            group.append(read_line)

        return Read(group[0], group[1], group[2], group[3])

    def close(self):
        self.file_hd.close()


async def send_chunks(file_path: str) -> str:
    """
    get some reads and send to streaming
    """
    fastq_iter = FastqIterator(file_path)
    n_bytes, read_list = 0, []

    while True:
        read_obj = fastq_iter.get_read()
        if not read_obj:  # all chunks have been sent
            fastq_iter.close()
            break

        read_list.append(read_obj.to_string())
        n_bytes += read_obj.read_length()

        if n_bytes >= CHUNK_SIZE:
            yield ''.join(read_list)
            n_bytes, read_list = 0, []

    if read_list:
        yield ''.join(read_list)


app = FastAPI(title="Stream Compute Server")


# register the tortoise-orm database
register_tortoise(app=app, config=TORTOISE_ORM, generate_schemas=True)


@app.get("/access", summary="get access to the target fastq file")
async def download_file(access_code: str):
    item_list = await PublicModel.filter(access_code=access_code).all()

    if not item_list:
        raise HTTPException(status_code=423, detail=F"Error: access_code of {access_code} is not exist!")

    item = item_list[0]

    # update the retrieve times of the data
    item.retrieve_times += 1
    await item.save()

    return StreamingResponse(send_chunks(item.file_path), media_type="text/plain")


@app.get("/list", summary="get public sample list")
async def get_public_list():
    """
    get all available public sample list
    """
    sample_list = await PublicModel.all().values(
        'access_code', 'species', 'retrieve_times', 'file_name'
    )

    return {
        'count': len(sample_list),
        'tasks': sample_list
    }


if __name__ == '__main__':
    uvicorn.run("server:app", host=SERVER_HOST, port=SERVER_PORT, reload=True)
